package com.lun.medium;

import java.util.LinkedList;

import com.lun.util.BinaryTree.TreeNode;

public class BinaryTreePruning {
	
	//方法一：用到后序遍历模式的递归版
    public TreeNode pruneTree(TreeNode root) {
        if(root == null) return null;
    	
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
    	
    	if(root.val == 0 && root.left == null && root.right == null)
    		return null;
    	
    	return root;
    }
    
    //方法二：用到后序遍历模式的迭代版
    public TreeNode pruneTree2(TreeNode root) {
    	if(root == null) return null;
    	
    	LinkedList<TreeNode[]> stack = new LinkedList<>();//需要将节点以及其父节点存入同一栈帧。
    	TreeNode fakeRoot = new TreeNode(-1, root, null);//对于根节点没有父节点的情况，可以给它弄个临时父节点。临时父节点的左节点指向根节点
    	TreeNode p = root, lastNodeVisited = null, parent = fakeRoot;
    	
    	while(!stack.isEmpty() || p != null) {
    		
    		if(p != null) {
    			stack.push(new TreeNode[] {parent, p});
    			parent = p;
    			p = p.left;
    		}else {
    			TreeNode[] peekOne = stack.peek();
    			if(peekOne[1].right != null && peekOne[1].right != lastNodeVisited) {
    				parent = peekOne[1];
    				p = peekOne[1].right;
    			}else {
    				if(peekOne[1].val == 0 && peekOne[1].left == null && peekOne[1].right == null) {
    					if(peekOne[0].left == peekOne[1]) {
    						peekOne[0].left = null;
    					}else {
    						peekOne[0].right = null;
    					}
    				}
    				lastNodeVisited = stack.pop()[1];
    			}
    		}
    	}
    	
    	return fakeRoot.left;
    }

}
